from operator import itemgetter

from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool

db = SQLDatabase.from_uri("sqlite:///../my_db.db")

# print(db.get_usable_table_names())
# print(db.run("SELECT * FROM tb_user"))


# 第一步：自然语言转SQL
key = 'sk-7wnDma9l5GuVbq38B3C07f50290147148a0809B117A1C1Ad'
model = ChatOpenAI(model="gpt-3.5-turbo",
                   openai_api_key=key,
                   openai_api_base="https://api.aigc369.com/v1",
                   temperature=0)

query_chain = create_sql_query_chain(db=db, llm=model)

# 第二步：SQL查询
sql_executor = QuerySQLDataBaseTool(db=db)

executor_chain = query_chain | sql_executor

# 第三步：SQL结果和原始问题一起交给模型

prompt_content = """
    你是一个数据库信息回答助手，你可以根据用户给的问题，待执行的SQL查询语句以及SQL查询之后的结果，进行一个总结输出给用户
    
    问题:{question}
    
    待执行的SQL查询语句:{sql}

    SQL查询之后的结果:{result}

    现在开始!
"""

prompt = ChatPromptTemplate.from_template(prompt_content)

chain = {"question": itemgetter("question"), "sql": query_chain, "result": executor_chain} | prompt | model

print(chain.invoke({"question": "有多少个用户"}))


# 新版agent


SQL_PREFIX = """You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

To start you should ALWAYS look at the tables in the database to see what you can query.
Do NOT skip this step.
Then you should query the schema of the most relevant tables."""

# 老版agent
agent_executor = create_sql_agent(
    llm=model,db=db, agent_type="openai-tools", verbose=True, handle_parsing_errors=True
)

print(agent_executor.invoke({"input": "张三属于哪个部门，我需要部门的名字"}))
